{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "52295407",
   "metadata": {},
   "outputs": [],
   "source": [
    "import logging\n",
    "import sys\n",
    "\n",
    "logging.basicConfig(stream=sys.stdout, level=logging.INFO)\n",
    "logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c5d167a5-81f8-4d2c-b42f-0a190577132f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# My OpenAI Key\n",
    "import os\n",
    "\n",
    "os.environ[\"OPENAI_API_KEY\"] = \"INSERT OPENAI KEY\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "575750cc-479f-4b1f-b93f-4b00ed756d52",
   "metadata": {},
   "source": [
    "## Wikipedia Reader + Keyword Table"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "5f60348e-731d-4a95-bae2-426e184a914e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index import KeywordTableIndex, WikipediaReader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "952c4659-7fbb-447e-8caf-06916412cc37",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "page: Covid-19\n"
     ]
    }
   ],
   "source": [
    "wiki_docs = WikipediaReader().load_data(pages=[\"Covid-19\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3be202db-a4c7-41d2-ba7d-446d1f934830",
   "metadata": {},
   "outputs": [],
   "source": [
    "index = KeywordTableIndex.from_documents(wiki_docs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "28d7163e-f26f-4ad8-89d5-9cb7662c4d9c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "> Starting query: Which country included tocilizumab in treatment for covid-19?\n",
      "Extracted keywords: ['tocilizumab', 'treatment', 'covid-19', 'covid', '19']\n",
      "> Querying with idx: 1105763466456338724: of age or older weighing at least 40 kilograms ...\n",
      "> Querying with idx: 2820318727532393752: Coronavirus disease 2019 (COVID-19) is a contag...\n",
      "> Querying with idx: 897499143815831368: if the mask includes an exhalation valve, a wea...\n",
      "> Querying with idx: 8628144746434065339: pulmonary fibrosis, cystic fibrosis. Evidence s...\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'\\n\\nChina'"
      ]
     },
     "execution_count": 46,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# GPT doesn't find the corresponding evidence in the leaf node, but still gives the correct answer\n",
    "# set Logging to DEBUG for more detailed outputs\n",
    "query_engine = index.as_query_engine()\n",
    "query_engine.query(\"Which country included tocilizumab in treatment for covid-19?\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "addb0c4d-f1ae-40c1-8b69-5a989609672f",
   "metadata": {},
   "source": [
    "## Wikipedia Reader + List"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a0fc24e1-eca5-4267-a962-f7fe0fc5c7df",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index import ListIndex, WikipediaReader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "872a651a-ca4a-43e2-8b29-e4f667f9d3c5",
   "metadata": {},
   "outputs": [],
   "source": [
    "wiki_docs = WikipediaReader().load_data(pages=[\"Covid-19\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "37e85af0-b1c3-4c18-b239-6e32a7acf8d6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "> Adding chunk: Coronavirus disease 2019 (COVID-19) is a contag...\n",
      "> Adding chunk: people with COVID‑19 and acute respiratory dist...\n",
      "> Adding chunk: encourage or mandate the use of face masks or c...\n",
      "> Adding chunk: have elevated liver enzymes, reflecting liver i...\n",
      "> Adding chunk: insofar as their drug use may have caused lung ...\n",
      "> Adding chunk: treatment of mild-to-moderate COVID‑19 in adult...\n"
     ]
    }
   ],
   "source": [
    "index = ListIndex.from_documents(wiki_docs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "ec0119ef-786e-40ea-89af-f1ca0ad26de6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "> Starting query: Which country included tocilizumab in treatment for covid-19?\n"
     ]
    }
   ],
   "source": [
    "# set Logging to DEBUG for more detailed outputs\n",
    "# with keyword lookup\n",
    "from llama_index.indices.postprocessor import KeywordNodePostprocessor\n",
    "\n",
    "\n",
    "query_engine = index.as_query_engine(\n",
    "    node_postprocessors=[KeywordNodePostprocessor(required_keywords=[\"tocilizumab\"])]\n",
    ")\n",
    "response = query_engine.query(\n",
    "    \"Which country included tocilizumab in treatment for covid-19?\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "b4087a84-0939-444f-93f2-a1a7aa32db3f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'China'"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "display(response.strip())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "fb155bc7-cb50-47b6-b92b-895852c2d8f4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "> Starting query: Which country included tocilizumab in treatment for covid-19?\n"
     ]
    }
   ],
   "source": [
    "# set Logging to DEBUG for more detailed outputs\n",
    "# without keyword lookup\n",
    "query_engine = index.as_query_engine()\n",
    "response = query_engine.query(\n",
    "    \"Which country included tocilizumab in treatment for covid-19?\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "5b45c07a-4e76-4a45-86b6-6b2df1ef4f7b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'There is no definite answer to this question as different countries have different treatment methods for covid-19. However, according to the context information, it is known that the virus SARS-CoV-2 can cause severe damage to various organs in the human body by inducing systemic inflammation. Therefore, it is possible that tocilizumab, which is a drug that inhibits the virus, may be included in treatment for covid-19 in some countries in order to prevent or reduce the severity of a cytokine storm. Additionally, passive antibodies may be used to treat people with active COVID-19 in order to help them recover.'"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "display(response.strip())"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
